from collections import Counter
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import degree
import numpy as np
import torch
import torch.nn as nn

# 利用PyG的消息传递接口来获取感染邻居数目
class i_neighbor(MessagePassing):
    def __init__(self):
        super().__init__(aggr="add")

    def forward(self, x, edge_index):
        return self.propagate(edge_index, x=x)


# 带权重的损失
class weighted_loss(nn.Module):
    def __init__(self, edge_index, lamb=0.5):
        super().__init__()

        self.edge_index = edge_index
        self.lamb = lamb
        self.deg = (
            degree(edge_index[0])
            .to(torch.long)
            
        )
        self.i_n = i_neighbor()

    def forward(self, out, data):
        x = data.x
        y = data.y.long()
        w = self.weight(x)
        w /= w.sum()

        out = torch.clamp(out, 1e-15, 1 - 1e-15)
        loss = (w * (-y * torch.log(out)).sum(-1)).sum()

        return loss

    def weight(self, x):
        x_state = x.view(-1).to(torch.long)
        infected_neighbors = self.i_n(x, self.edge_index).view(-1).to(torch.long)

        config = list(
            map(
                lambda x: tuple(x),
                (
                    torch.stack((self.deg, x_state, infected_neighbors), dim=1)
                    .to("cpu")
                    .numpy()
                    .tolist()
                ),
            )
        )

        pr = Counter(config)
        w = np.zeros(len(x))
        for i, item in enumerate(config):
            w[i] = pow(float(pr[item] / len(x)), -self.lamb)

        return torch.tensor(w)
